# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/21 12:34 
@Author  : ZhangShenao 
@File    : app_handler.py
@Desc    : 应用控制器
"""
import uuid
from dataclasses import dataclass

from injector import inject
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

from internal.schema import ChatReq
from internal.service import AppService, VectorStoreService
from pkg.response import validate_error, success, success_with_message
from .chat_memory_chain import invoke_chain_with_chat_memory, INPUT_KEY, HISTORY_KEY

# 系统prompt
SYSTEM_PROMPT = '''你是一个强大的聊天机器人，请根据用户的当前提问，结合历史聊天记录和上下文信息，做出准确的回答。
下面是相关的上下文信息:
<context>
{context}
</context>
'''

# 模型名称
MODEL_NAME = 'GLM-4-Air'


@inject  # 设置依赖注入
@dataclass
class AppHandler:
    """
    应用控制器
    """

    app_service: AppService  # 依赖注入AppService

    vector_store_service: VectorStoreService  # 依赖注入VectorStoreService

    def debug(self):
        """
        应用聊天调试
        """

        # 请求参数校验
        req = ChatReq()
        if not req.validate():
            return validate_error(req.errors)
        input_data = req.query.data

        # 构造Prompt、LLM、OutputParser和Memory等组件
        prompt = ChatPromptTemplate.from_messages([
            ('system', SYSTEM_PROMPT),
            MessagesPlaceholder(HISTORY_KEY),  # 使用MessagesPlaceholder消息占位符
            ('human', '{input}')
        ])
        llm = ChatZhipuAI(model=MODEL_NAME, temperature=0.5)
        parser = StrOutputParser()

        # 调用封装好的Memory Chain,获取结果
        output = invoke_chain_with_chat_memory(input={INPUT_KEY: input_data},
                                               prompt_template=prompt,
                                               llm=llm,
                                               parser=parser,
                                               vector_store_service=self.vector_store_service)

        # 返回结果
        return success({'content': output})

    def create(self):
        """
        创建AI应用
        """

        app = self.app_service.create()
        return success_with_message(f'新建AI应用成功，应用ID为: {app.id}')

    def get_by_id(self, app_id: uuid.UUID):
        app = self.app_service.get_by_id(app_id)
        return success_with_message(f'查询AI应用成功，名称为: {app.name}')

    def get_all(self):
        apps = self.app_service.get_all()
        return success_with_message(f'查询所有AI应用成功，总数为: {len(apps)}')

    def update(self, app_id: uuid.UUID):
        app = self.app_service.update(app_id)
        return success_with_message(f'更新AI应用成功，应用ID为: {app.id}')

    def delete(self, app_id: uuid.UUID):
        app = self.app_service.delete(app_id)
        return success_with_message(f'删除AI应用成功，应用ID为: {app.id}')
